"""
联系人表结构迁移脚本
"""

from sqlalchemy import create_engine, text
from config import settings

def migrate_contacts_table():
    """迁移联系人表结构"""
    engine = create_engine(settings.database_url)
    
    with engine.connect() as conn:
        # 开始事务
        trans = conn.begin()
        
        try:
            # 检查contacts表是否存在
            result = conn.execute(text("""
                SELECT COUNT(*) 
                FROM information_schema.tables 
                WHERE table_schema = DATABASE() 
                AND table_name = 'contacts'
            """))
            
            table_exists = result.scalar() > 0
            
            if table_exists:
                print("联系人表已存在，检查是否需要更新字段...")
                
                # 检查是否有user_id字段
                result = conn.execute(text("""
                    SELECT COUNT(*) 
                    FROM information_schema.columns 
                    WHERE table_schema = DATABASE() 
                    AND table_name = 'contacts' 
                    AND column_name = 'user_id'
                """))
                
                has_user_id = result.scalar() > 0
                
                if not has_user_id:
                    print("添加user_id字段...")
                    conn.execute(text("""
                        ALTER TABLE contacts 
                        ADD COLUMN user_id INT NOT NULL DEFAULT 1,
                        ADD COLUMN company VARCHAR(200),
                        ADD COLUMN position VARCHAR(100),
                        MODIFY COLUMN email VARCHAR(255) NOT NULL,
                        MODIFY COLUMN phone VARCHAR(50)
                    """))
                    
                    # 添加外键约束
                    conn.execute(text("""
                        ALTER TABLE contacts 
                        ADD CONSTRAINT fk_contacts_user_id 
                        FOREIGN KEY (user_id) REFERENCES users(id)
                    """))
                    
                    # 添加索引
                    conn.execute(text("""
                        CREATE INDEX idx_contacts_user_email ON contacts(user_id, email)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_contacts_user_name ON contacts(user_id, name)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_contacts_company ON contacts(company)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_contacts_created_at ON contacts(created_at)
                    """))
                    
                    print("联系人表字段更新完成")
                else:
                    print("联系人表字段已是最新")
            else:
                print("联系人表不存在，将由SQLAlchemy自动创建")
            
            # 创建联系人分组表
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS contact_groups (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    name VARCHAR(100) NOT NULL COMMENT '分组名称',
                    description TEXT COMMENT '分组描述',
                    user_id INT NOT NULL COMMENT '用户ID',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                    FOREIGN KEY (user_id) REFERENCES users(id),
                    INDEX idx_contact_groups_user_name (user_id, name)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='联系人分组表'
            """))
            
            # 创建联系人分组成员关联表
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS contact_group_members (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    contact_id INT NOT NULL COMMENT '联系人ID',
                    group_id INT NOT NULL COMMENT '分组ID',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    FOREIGN KEY (contact_id) REFERENCES contacts(id) ON DELETE CASCADE,
                    FOREIGN KEY (group_id) REFERENCES contact_groups(id) ON DELETE CASCADE,
                    UNIQUE KEY idx_contact_group_members_unique (contact_id, group_id),
                    INDEX idx_contact_group_members_contact (contact_id),
                    INDEX idx_contact_group_members_group (group_id)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='联系人分组成员关联表'
            """))
            
            # 提交事务
            trans.commit()
            print("联系人表结构迁移完成！")
            
        except Exception as e:
            # 回滚事务
            trans.rollback()
            print(f"迁移失败: {e}")
            raise

if __name__ == "__main__":
    migrate_contacts_table()
